from pathlib import Path
from numpy import random
import time
import numpy as np
from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, \
    scale_coords, set_logging, increment_path
from utils.plots import plot_one_box
from utils.torch_utils import select_device, time_synchronized
import cv2 as cv
import torch


class ObjectDetector():
    def __init__(self, weights, imgsz, conf_thres, iou_thres, device):
        # Directories
        self.save_dir = Path(increment_path(
            Path(""), exist_ok=False))  # increment run
        self.weights = weights

        # 初始化
        set_logging()
        self.device = select_device(device)
        self.half = self.device.type != 'cpu'  # half precision only supported on CUDA

        # 加载模型
        self.model = attempt_load(
            weights, map_location=self.device)  # load FP32 model
        stride = int(self.model.stride.max())  # model stride
        self.imgsz = check_img_size(imgsz, s=stride)  # check img_size
        if self.half:
            self.model.half()  # to FP16

        # 获取模型层和随机颜色
        self.names = self.model.module.names if hasattr(
            self.model, 'module') else self.model.names
        self.colors = [[random.randint(0, 255)
                        for _ in range(3)] for _ in self.names]

        # 执行推理
        if self.device.type != 'cpu':
            self.model(torch.zeros(1, 3, self.imgsz, self.imgsz).to(
                self.device).type_as(next(self.model.parameters())))  # run once
        t0 = time.time()
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.output = ''

    def inference(self, img):
        """
        inference 
        部署模型推理结果

        Args:
            img:源图像

        Returns:
            im0:推理后图像
        """
        copeImg = letterbox(img, self.imgsz, stride=32)[0]
        # Convert
        # BGR to RGB, to 3x416x416
        copeImg = copeImg[:, :, ::-1].transpose(2, 0, 1)
        copeImg = np.ascontiguousarray(copeImg)
        copeImg = torch.from_numpy(copeImg).to(self.device)
        # uint8 -> fp16/32
        copeImg = copeImg.half() if self.half else copeImg.float()
        copeImg /= 255.0
        if copeImg.ndimension() == 3:
            copeImg = copeImg.unsqueeze(0)
        # 推理
        t1 = time_synchronized()
        pred = self.model(copeImg, False)[0]
        # 非极大值抑制
        pred = non_max_suppression(
            pred, self.conf_thres, self.iou_thres, classes=None, agnostic=True)
        t2 = time_synchronized()

        isfind = False
        result = []
        # 处理推理结果
        for i, det in enumerate(pred):  # detections per image
            s, im0 = '', img
            if len(det):
                # 将得到的尺寸映射回源图像尺寸
                det[:, :4] = scale_coords(
                    copeImg.shape[2:], det[:, :4], im0.shape).round()

                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    # add to string
                    s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "

                # 得到结果并绘制
                for *xyxy, conf, cls in reversed(det):
                    label = f'{self.names[int(cls)]} {conf:.2f}'
                    if self.names[int(cls)] == 'person':
                        isfind = True
                        result.append(xyxy)
                        plot_one_box(xyxy, im0, label=label,
                                     color=self.colors[int(cls)], line_thickness=3)
            # 打印输出和NMS
            # self.output = f'{s}Done. ({t2 - t1:.3f}s)'
            self.output = f'{s}.  '
            if not isfind:
                self.output = "无目标  "
            # print(self.output)
            return im0, result


def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    shape = img.shape[:2]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:
        r = min(r, 1.0)

    # 计算Padding
    ratio = r, r
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - \
        new_unpad[1]
    if auto:
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / \
            shape[0]

    dw /= 2
    dh /= 2

    if shape[::-1] != new_unpad:
        img = cv.resize(img, new_unpad, interpolation=cv.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv.copyMakeBorder(img, top, bottom, left, right,
                            cv.BORDER_CONSTANT, value=color)
    return img, ratio, (dw, dh)
